#!/usr/bin/env python
from astropy.io import fits
from astropy.table import Table
import matplotlib.pyplot as plt
import numpy as np
import sys

din = sys.argv[1]  #暗场处理后文件夹
dbad = sys.argv[2] #坏像素处理后文件夹
tblock = sys.argv[3]  #需要遮挡温像元的温度，多个温度用逗号分隔
ts = tblock.split(',')

temp = []
for t in ts:
    temp.append(int(t))
    
for t in temp:
    print("processing temp %i for %s,%s"%(t,din,dbad))
    
    print("reading input files")
    outfile = din+"/dark_outlier"+str(t)+"k_08.fits"
    print("outfile",outfile)
    dark_tab = Table.read(din+"/dark_stats.tab",format='ipac')
    npix08 = dark_tab['npix08'][dark_tab['tsys'] == t]
    ratio08 = dark_tab['frac08'][dark_tab['tsys'] == t]
    print("before blocking",npix08[0],ratio08[0])
    outmap = fits.getdata(outfile)
    
    print("searching warm columns")
    warm_cols = []
    for i in range(outmap.shape[1]):
        outcol = outmap[:,i]
        idxs = np.where(outcol==1)[0]
        #if len(idxs)>10:
        #    print("col",i+1,"idxs",idxs)
        count = 0
        for j in range(len(idxs)-1):
            if idxs[j+1]-idxs[j]<10:
                count += 1
                if count>=100:
                    break
            else:
                count = 0
        #if count>20:
            #print("col",i+1,"count",count)
            #print(idxs[:100])
            #plt.plot(outcol)
            #plt.show()
        if count==100:
            warm_cols.append(i)
            npix08 -= count
            ratio08 -= count/9232/9216
            
    
    print("modifying dark table")
    dark_tab['npix08'][dark_tab['tsys'] == t] = npix08
    dark_tab['frac08'][dark_tab['tsys'] == t] = ratio08
    print("after blocking",npix08[0],ratio08[0])
    dark_tab.write(din+"/dark_stats_block.tab",format='ipac',overwrite=True)

    print("modifying bad fits")
    bad = fits.getdata(dbad+"/bad.fits")
    for warm_col in warm_cols:
        bad[:,warm_col] = 2
    bad[:2] = 0
    bad[-2:] = 0
    bad[:,:2] = 0
    bad[:,-2:] = 0
    fits.writeto(dbad+"/bad_warm_block.fits",bad,overwrite=True)
    badcol = np.sum(bad[3] == 2)
    
    print("modifying bad table")
    bad_tab = Table.read(dbad+"/bad_stats.tab",format='ipac')
    bad_tab['col'] = [badcol/9216]
    bad_tab['lines'] = max([bad_tab['col'][0],bad_tab['row'][0]])
    bad_tab.write(dbad+"/bad_stats_warm_block.tab",format='ipac',overwrite=True)

